#include <thrust/copy.h>
#include <cusp/precond/bspai.h>
#include <cusp/csr_matrix.h>
#include <cusp/io/matrix_market.h>
#include <iostream>
#include <cusp/frobenius_norm.h>
#include <cusp/multiply.h>
#include <cusp/gallery/eye.h>
#include <cusp/elementwise.h>
#include <assert.h>


int main(int argc, char** argv)
{
    typedef unsigned int        IndexType;
    typedef unsigned int        IndexType2;
    typedef float               ValueType;
    typedef cusp::device_memory MemorySpace;


    std::cout << "Loading..." << std::endl;
    // create an empty sparse matrix structure
    cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> A;
    cusp::csr_matrix<IndexType, IndexType2, cusp::host_memory> A_BSP;
    
    cusp::io::read_matrix_market_file(A, std::string(argv[1]));
    cusp::io::read_matrix_market_file(A_BSP, std::string(argv[2]));
    
    std::cout << "Matrix size: " << A.num_rows << " with " << A.num_entries << " non-zeros." << std::endl;
    std::cout << "Block sparsity pattern size: " << A_BSP.num_rows << " with " << A_BSP.num_entries << " non-zero blocks." << std::endl;

    cusp::csr_matrix<IndexType, ValueType, MemorySpace> d_A(A);
    cusp::csr_matrix<IndexType, IndexType2, MemorySpace> d_A_BSP(A_BSP);
    
    std::cout << "Finished loading" << std::endl;

    std::cout << "Starting preconditioner" << std::endl;
    try
    {
        cusp::precond::bspai<IndexType, ValueType, MemorySpace> d_M(d_A, d_A_BSP);
   

        
        d_M.matrix.sort_by_row(); // XXX
        std::cout << "Completed preconditioner" << std::endl;
        
        /*
       	cusp::csr_matrix<IndexType, ValueType, MemorySpace> d_I;
        cusp::gallery::eye(d_I, A.num_rows, A.num_cols);
        cusp::csr_matrix<IndexType,ValueType, MemorySpace> d_P;
        cusp::multiply(d_M.matrix, d_A, d_P);
        cusp::subtract(d_P, d_I, d_P);
        std::cout << "#Fnorm(A)\t" << cusp::frobenius_norm(d_A) << "\n";
        std::cout << "#Fnorm(MA-I)\t" << cusp::frobenius_norm(d_P) << std::endl;
        */
        
        std::string newFname;
        std::string fname = std::string(argv[2]);
        std::string ext(".mtx");
        
        newFname=fname.substr(0, fname.find(ext)).append("_BSPAI.mtx");
        cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> M(d_M.matrix);
        cusp::io::write_matrix_market_file(M, newFname);
    }
    catch(const cusp::precond::detail::device::bspaiExcept& err)
    {
        std::cerr << err.getMsg(); 
        std::cout << err.getMsg();
        return 1;   
    }
    
    return 0;
}



